#include "pch.h"
#include "PuyaDfu.h"

#define READ_DATA_LENGTH	0x40

#define ERROR_COMMAND                     0xEC             /* Error command */
#define ACK_BYTE                          0x79             /* Acknowledge Byte ID */
#define NACK_BYTE                         0x1F             /* No Acknowledge Byte ID */
#define BUSY_BYTE                         0x76             /* Busy Byte */
#define SYNC_BYTE                         0xA5             /* synchronization byte */

#define CMD_GET_COMMAND                   0x00             /* Get commands command */
#define CMD_GET_VERSION                   0x01             /* Get Version command */
#define CMD_GET_ID                        0x02             /* Get ID command */
#define CMD_GET_DEVICE_IDCODE             0x03             /* Get DEVICE_IDCODE command */
#define CMD_READ_MEMORY                   0x11             /* Read Memory command */
#define CMD_WRITE_MEMORY                  0x31             /* Write Memory command */
#define CMD_GO                            0x21             /* GO command */
#define CMD_READ_PROTECT                  0x82             /* Readout Protect command */
#define CMD_READ_UNPROTECT                0x92             /* Readout Unprotect command */
#define CMD_EXT_ERASE_MEMORY              0x44             /* Erase Memory command */
#define CMD_WRITE_PROTECT                 0x63             /* Write Protect command */
#define CMD_WRITE_UNPROTECT               0x73             /* Write Unprotect command */
#define CMD_WRITE_PERIPHERALS             0xF1
#define CMD_READ_PERIPHERALS			  0xF2
#define CMD_SOFT_RESET					  0xF3

//unsigned char USB_WRITE = USB_WRITE;

bool CPuyaDfu::WaitData(const uint8_t ucData)
{
	uint8_t ucRead[0x40] = { 0x00 };

	for (int i = 0; i < 3; i++) {
		if (READ_DATA_LENGTH != usb_interrupt_transfer(USB_READ, ucRead, 0x40, 1000)) {
			TRACE(_T("WatiData Error.\n"));
			return false;
		}
		if (ucData == ucRead[0]) {
			return true;
		}
	}

	return false;
}

uint32_t CPuyaDfu::Init(uint32_t nPort, uint32_t nPID)
{
	int nErrCode = ERROR_SUCCESS;
	uint8_t ucData[0x40] = { 0x00 };

	do
	{
		if (0 >= usb_get_device_list(VID, nPID)) {
			return ERROR_USB_GET_DEVICE_LIST;
		}
		nErrCode = usb_open(0, nPort);
		if (ERROR_SUCCESS != nErrCode) {
			return nErrCode;
		}

		for (UINT i = 0; i < 2; i++) {
			memset(ucData, 0x7F, 0x40);
			usb_interrupt_transfer(USB_WRITE, ucData, 0x40);

			memset(ucData, 0x00, 0x40);
			if (READ_DATA_LENGTH == usb_interrupt_transfer(USB_READ, ucData, 0x40, 1000)) {
				nErrCode = ERROR_SUCCESS;
				break;
			}
			TRACE(_T("WatiData Error.\n"));
			nErrCode = ERROR_ISP_0x7F;
			//USB_WRITE = USB_WRITE - 1;
		}

		if (ERROR_SUCCESS != nErrCode) {
			//USB_WRITE = USB_WRITE;
			return nErrCode;
		}

		if ((ACK_BYTE != ucData[0]) && (NACK_BYTE != ucData[0])) {
			return ERROR_ISP_0x7F;
		}

		if (!Get()) {
			return ERROR_ISP_Get;
		}

		if (!GetID()) {
			return ERROR_ISP_GetID;
		}

		if (ERROR_SUCCESS == ReadMemory(0x20000800, ucData, 0x04))
		{
			return ERROR_SUCCESS;
		}

		if (IDNO == AfxMessageBox(_T("MCU read protection active, do you want to remove read protection?"), MB_YESNO|MB_ICONQUESTION))
		{
			return ERROR_IAP_READ;
		}

		if (!ReadoutUnprotect())
		{
			return ERROR_ISP_ReadoutUnprotect;
		}

		usb_close();

		Sleep(500);

	} while (1);

	return ERROR_SUCCESS;
}

uint32_t CPuyaDfu::Uninit(void)
{
	usb_close();
	return ERROR_SUCCESS;
}

uint32_t CPuyaDfu::WriteMemory(uint32_t addr, uint8_t* data, uint32_t size)
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = CMD_WRITE_MEMORY;
	ucData[0x01] = 0xFF - ucData[0x00];

	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return ERROR_IAP_WRITE;
	}
	ucData[0x00] = HIBYTE(HIWORD(addr));
	ucData[0x01] = LOBYTE(HIWORD(addr));
	ucData[0x02] = HIBYTE(LOWORD(addr));
	ucData[0x03] = LOBYTE(LOWORD(addr));
	ucData[0x04] = GetXOR(ucData, 0x04);
	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return ERROR_IAP_WRITE;
	}

	memset(ucData, 0x00, 0x40);
	ucData[0x00] = size - 1;
	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);

	uint32_t i;
	for (i = 0; i < size / 0x40; i++) {
		usb_interrupt_transfer(USB_WRITE, &data[i * 0x40], 0x40);
	}
	if (size % 0x40) {
		memset(ucData, 0x00, 0x40);
		memcpy(ucData, &data[i * 0x40], size % 0x40);
		ucData[size % 0x40] = GetXOR(data, size, size - 1);//CRC
		usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	}
	else {
		memset(ucData, 0x00, 0x40);
		ucData[0] = GetXOR(data, size, size - 1);//CRC
		usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	}

	if (!WaitData(ACK_BYTE)) {
		return ERROR_IAP_WRITE;
	}

	return ERROR_SUCCESS;
}

uint8_t  CPuyaDfu::ReadMemory(uint32_t addr, uint8_t* data, uint32_t size)
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = CMD_READ_MEMORY;
	ucData[0x01] = 0xFF - ucData[0x00];

	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return ERROR_IAP_READ;
	}

	ucData[0x00] = HIBYTE(HIWORD(addr));
	ucData[0x01] = LOBYTE(HIWORD(addr));
	ucData[0x02] = HIBYTE(LOWORD(addr));
	ucData[0x03] = LOBYTE(LOWORD(addr));
	ucData[0x04] = GetXOR(ucData, 0x04);
	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return ERROR_IAP_READ;
	}

	memset(ucData, 0x00, 0x40);
	ucData[0x00] = size - 1;
	ucData[0x01] = 0xFF - ucData[0x00];
	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return ERROR_IAP_READ;
	}

	uint32_t i;
	for (i = 0; i < size / 0x40; i++) {
		if (READ_DATA_LENGTH != usb_interrupt_transfer(USB_READ, &data[i * 0x40], 0x40, 1000))
		{
			return ERROR_IAP_READ;
		}
	}
	if (size % 0x40) {
		if (READ_DATA_LENGTH != usb_interrupt_transfer(USB_READ, ucData, 0x40, 1000))
		{
			return ERROR_IAP_READ;
		}
		memcpy(&data[i * 0x40], ucData, size % 0x40);
	}

	return ERROR_SUCCESS;
}

uint32_t CPuyaDfu::ErasePage(uint8_t* data, uint32_t size)
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = 0x10;
	ucData[0x01] = size - 1;

	memcpy(ucData + 2, data, size * 2);

	if (!ExtendedErase(ucData, (0x01 + size) * 2))
	{
		return ERROR_IAP_ERASE_PAGE;
	}
	return ERROR_SUCCESS;
}

uint32_t CPuyaDfu::EraseSector(uint8_t* data, uint32_t size)
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = 0x20;
	ucData[0x01] = size - 1;

	memcpy(ucData + 2, data, (size_t)(size * 2));

	if (!ExtendedErase(ucData, (0x01 + size) * 2))
	{
		return ERROR_IAP_ERASE_SECTOR;
	}
	return ERROR_SUCCESS;
}

uint32_t CPuyaDfu::EraseChip(void)
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = 0xFF;
	ucData[0x01] = 0xFF;

	if (!ExtendedErase(ucData, 0x02))
	{
		return ERROR_IAP_ERASE_ALL;
	}

	return ERROR_SUCCESS;
}

uint8_t CPuyaDfu::Go(uint32_t addr)
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = CMD_GO;
	ucData[0x01] = 0xFF - ucData[0x00];

	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return ERROR_IAP_RUN_APP;
	}

	ucData[0x00] = HIBYTE(HIWORD(addr));
	ucData[0x01] = LOBYTE(HIWORD(addr));
	ucData[0x02] = HIBYTE(LOWORD(addr));
	ucData[0x03] = LOBYTE(LOWORD(addr));
	ucData[0x04] = GetXOR(ucData, 0x04);
	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return ERROR_IAP_RUN_APP;
	}

	return ERROR_SUCCESS;
}

uint32_t CPuyaDfu::WriteOptionBytes(uint32_t addr, uint8_t* data, uint32_t size)
{
	uint32_t nErrCode;

	nErrCode = WriteMemory(addr, data, size);
	if (ERROR_SUCCESS != nErrCode) {
		return nErrCode;
	}

	return ERROR_SUCCESS;
}

bool CPuyaDfu::Get()
{
	int nErrCode = 0;
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = CMD_GET_COMMAND;
	ucData[0x01] = 0xFF - ucData[0x00];

	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	memset(ucData, 0x00, 0x40);
	usb_interrupt_transfer(USB_READ, ucData, 0x40);

	return true;
}

bool CPuyaDfu::GetVersionAndReadProtectionStatus()
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = CMD_GET_VERSION;
	ucData[0x01] = 0xFF - ucData[0x00];

	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);

	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	memset(ucData, 0x00, 0x40);
	usb_interrupt_transfer(USB_READ, ucData, 0x40);

	return true;
}

bool CPuyaDfu::GetID()
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = CMD_GET_ID;
	ucData[0x01] = 0xFF - ucData[0x00];

	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);

	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	memset(ucData, 0x00, 0x40);
	usb_interrupt_transfer(USB_READ, ucData, 0x40);

	//m_wPID = MAKEWORD(ucData[0x02], ucData[0x01]);

	return true;
}

bool CPuyaDfu::ExtendedErase(BYTE* pucData, BYTE ucSize)
{
	BYTE ucData[0x40] = { 0x00 };

	ucData[0x00] = CMD_EXT_ERASE_MEMORY;
	ucData[0x01] = 0xFF - ucData[0x00];
	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	memcpy(ucData, pucData, 2);
	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);

	memset(ucData, 0x00, 0x40);
	memcpy(ucData, pucData + 2, ucSize - 2);
	ucData[ucSize - 2] = GetXOR(pucData, ucSize, 0x00);
	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	return true;
}

bool CPuyaDfu::WriteProtect(BYTE* pucData, BYTE ucSize)
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = CMD_WRITE_PROTECT;
	ucData[0x01] = 0xFF - ucData[0x00];

	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	ucData[0x00] = ucSize - 1;
	memcpy(&ucData[0x01], pucData, ucSize);
	ucData[ucSize + 1] = GetXOR(ucData, ucSize + 1);
	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	return true;
}

bool CPuyaDfu::WriteUnprotect()
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = CMD_WRITE_UNPROTECT;
	ucData[0x01] = 0xFF - ucData[0x00];

	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	return true;
}

bool CPuyaDfu::ReadoutProtect()
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = CMD_READ_PROTECT;
	ucData[0x01] = 0xFF - ucData[0x00];

	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	return true;
}

bool CPuyaDfu::ReadoutUnprotect()
{
	uint8_t ucData[0x40] = { 0x00 };

	ucData[0x00] = CMD_READ_UNPROTECT;
	ucData[0x01] = 0xFF - ucData[0x00];

	usb_interrupt_transfer(USB_WRITE, ucData, 0x40);
	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	if (!WaitData(ACK_BYTE)) {
		return false;
	}

	return true;
}